import os, json, pprint
import itertools
from typing import Iterator, Optional

import streamlit as st

import newapi
from newapi import generate_roles_from_description, combine_roles,generate_role_play
from data_types import TextMsg, ImageMsg, TextMsgList, MsgList, filter_text_msg

st.set_page_config(page_title="Role-Play", page_icon="🤖", layout="wide")

def update_api_key(key: Optional[str] = None):
    # if debug:
    #     print(f'update_api_key. st.session_state["API_KEY"] = {st.session_state["API_KEY"]}, key = {key}')
    key = key or st.session_state["API_KEY"]
    if key:
        newapi.API_KEY = key

# 设置API KEY
api_key = st.sidebar.text_input("API_KEY", value=os.getenv("API_KEY", ""), key="API_KEY", type="password", on_change=update_api_key)
update_api_key(api_key)

# 初始化
if "description" not in st.session_state:
    st.session_state["description"] = {
        "text": ""
    }

# description_labels = {
#     "text": "人设描述"
# }

# def on_click_generate():
#     # 生成人设
#     print(st.session_state["description"]['text'])
#     st.markdown(st.session_state["description"]['text'])

# 生成人设
with st.container():
    col1,col2 = st.columns(2)
    with col1:
        # st.text_area(label = "人设描述", key = "description_text", on_change = lambda: st.session_state["description"].update(text = st.session_state["description_text"]), help = "输入文本")
        role_description_1 = st.text_area(label = "人设1描述", key = "description_text1", help = "输入文本")
        role_description_2 = st.text_area(label = "人设2描述", key = "description_text2", help = "输入文本")
        bt1, bt2 = st.columns(2)
        with bt1:
            generate_roles = st.button("生成人设")
        with bt2:
            clear_roles = st.button("清空人设")

    with col2:
        topic_description = st.text_area(label = "话题描述", key = "topic_description", help = "输入文本")
        chat_turns = st.number_input(label = "对话轮数", key = "chat_turns", value = 10, min_value = 1, max_value = 20)
        bt1, bt2, bt3 = st.columns(3)
        with bt1:
            generate_chats = st.button("生成对话")   
        with bt2:
            clear_message = st.button("清空对话") 
        with bt3:
            save_message = st.button("保存对话")

if 'role1_system' not in st.session_state:
    st.session_state['role1_system'] = []
if 'role2_system' not in st.session_state:
    st.session_state['role2_system'] = []
# 初始化会话状态中的 role1 和 role2
# if 'role1' not in st.session_state:
#     st.session_state['role1'] = ""
# if 'role2' not in st.session_state:
#     st.session_state['role2'] = ""
if 'system' not in st.session_state:
    st.session_state['system'] = []
if generate_roles:
    # 生成人设
    if role_description_1 and role_description_2:
        text1 = "".join(generate_roles_from_description(role_description_1))
        text2 = "".join(generate_roles_from_description(role_description_2))
        # st.session_state["role1_system"] = "".join(generate_roles_from_description(role_description_1))
        # st.session_state["role2_system"] = "".join(generate_roles_from_description(role_description_2))
        st.session_state["system"] = combine_roles(text1, text2)
        print(f'combine text is {st.session_state["system"]}')
        # st.markdown(text1)
        # st.markdown(text2)
        st.session_state["role1_system"] = st.session_state['system'].split('|')[0]
        st.session_state["role2_system"] = st.session_state['system'].split('|')[1]
    else:
        st.error("人设描述不能为空")

if st.session_state['role1_system']:
    text = st.session_state['role1_system']
    print(text)
    with st.chat_message(text.split('：')[0]):
        st.text(text.split('：')[1])
    

if st.session_state['role2_system']:
    text = st.session_state['role2_system']
    print(text)
    with st.chat_message(text.split('：')[0]):
        st.text(text.split('：')[1])
#     print(st.session_state['role2_system'])
#     with st.chat_message(f"{st.session_state['system'].split('\n')[1]}):
#         st.markdown(st.session_state['system'].split('\n')[1])

    # print(f'text1 = {st.session_state["text"]}')

if 'role1_messages' not in st.session_state:
    st.session_state['role1_messages'] = []
if 'role2_messages' not in st.session_state:
    st.session_state['role2_messages'] = []
if 'dialog' not in st.session_state:
    st.session_state['dialog'] = []

if generate_chats:
    # 生成对话
    
    if st.session_state["role1_system"] and st.session_state["role2_system"]:
        topic = topic_description
        for i in range(chat_turns):
            response = generate_role_play(st.session_state["role1_system"],st.session_state.role1_messages,topic)
            if topic:
                with st.chat_message(name=f"{st.session_state['role2_system'].split('：')[0]}"):
                    st.markdown(f"{topic}")
                    st.session_state['role2_messages'].append(TextMsg({"role": "assistant", "content": topic}) )
                    st.session_state['role1_messages'].append(TextMsg({"role": "user", "content": topic}) )  
                    st.session_state['dialog'].append(st.session_state['role2_system'].split('：')[0] + '：' + topic)
                topic = ""
            print(response.choices[0].message.content)
            msg = response.choices[0].message         
            if msg.role == "assistant":
                
                with st.chat_message(name=f"{st.session_state['role1_system'].split('：')[0]}"):
                    st.markdown(f"{msg.content}")
                    st.session_state['role1_messages'].append(TextMsg({"role": "assistant", "content": msg.content}) )
                    st.session_state['role2_messages'].append(TextMsg({"role": "user", "content": msg.content}) )  
                    st.session_state['dialog'].append(st.session_state['role1_system'].split('：')[0] + '：' + msg.content)
            else:
                st.markdown(f"**{msg.role}**：{msg.content}")
            response = generate_role_play(st.session_state["role2_system"],st.session_state.role2_messages,topic)
            
            msg = response.choices[0].message         
            if msg.role == "assistant":
                with st.chat_message(name=f"{st.session_state['role2_system'].split('：')[0]}"):
                    st.markdown(f"{msg.content}")
                    st.session_state['role1_messages'].append(TextMsg({"role": "user", "content": msg.content}) )
                    st.session_state['role2_messages'].append(TextMsg({"role": "assistant", "content": msg.content}) ) 
                    st.session_state['dialog'].append(st.session_state['role2_system'].split('：')[0] + '：' + msg.content)
            else:
                st.markdown(f"**{msg.role}**：{msg.content}")
    else:
        st.error("人设不能为空")
    # else:
    #     st.error("话题描述不能为空")

if clear_roles:
    st.session_state["role1_system"] = []
    st.session_state["role2_system"] = []
    st.session_state["system"] = []

if clear_message:
    st.session_state["role1_messages"] = []
    st.session_state["role2_messages"] = []
    st.session_state["dialog"] = []

if save_message:
    pprint.pprint(st.session_state['dialog'])
    with open('dialog.txt', 'w', encoding='utf-8') as f:
        for item in st.session_state['dialog']:
            f.write("%s\n" % item)
    st.success("对话已保存")
    
    with open('dialog.txt', 'rb') as f:
        st.download_button(
            label="点击下载",
            data=f,
            file_name='dialog.txt',
            mime='text/plain'
        )